#include "mpi.h"
#include "oom.h"
#include "cm.h"
#include "env.h"
#include "dtypes.h"
#include "util.h"

static void *user_buffer = NULL;
static int user_buffer_size = 0;

int MPI_Send(void *buf, int count, MPI_Datatype datatype, int dest,
	int tag, MPI_Comm comm)
{
	int ierr;
	MPI_Request request;

	ierr = MPI_Isend(buf, count, datatype, dest, tag, comm, &request);
	if (ierr != MPI_SUCCESS) return ierr;

	return MPI_Wait(&request, MPI_STATUS_IGNORE);
}

int MPI_Recv(void *buf, int count, MPI_Datatype datatype, int source,
	int tag, MPI_Comm comm, MPI_Status *status)
{
	int ierr;
	MPI_Request request;
	
	ierr = MPI_Irecv(buf, count, datatype, source, tag, comm, &request);
	if (ierr != MPI_SUCCESS) return ierr;
	return MPI_Wait(&request, status);
}

int MPI_Get_count(MPI_Status *status, MPI_Datatype datatype, int *count)
{
	int ierr, type_size;
	
	ierr = MPI_Type_size(datatype, &type_size);
	if (ierr != MPI_SUCCESS) return MPI_ERR_TYPE;
	
	if (status->bytes_recvd % type_size > 0)
	{
		*count = MPI_UNDEFINED;
	}
	else
	{
		*count = status->bytes_recvd / type_size;
	}
	
	return MPI_SUCCESS;
}

int MPI_Bsend(void *buf, int count, MPI_Datatype datatype, int dest,
	int tag, MPI_Comm comm)
{
	return MPI_ERR_OTHER;
}

int MPI_Rsend(void *buf, int count, MPI_Datatype datatype, int dest,
	int tag, MPI_Comm comm)
{
	int rc;
	rc = MPI_Send(buf, count, datatype, dest, tag, comm);
	return rc;
}

int MPI_Ssend(void *buf, int count, MPI_Datatype datatype, int dest,
	int tag, MPI_Comm comm)
{
	int rc;
	rc = MPI_Send(buf, count, datatype, dest, tag, comm);
	return rc;
}

int MPI_Buffer_attach(void *buffer, int size)
{
	if (buffer == NULL) return MPI_ERR_BUFFER;
	user_buffer = buffer;
	user_buffer_size = size;
	return MPI_SUCCESS;
}

int MPI_Buffer_detach(void *buffer, int *size)
{
	/* FIXME: must wait for all pending requests using buffer to complete */
	*((void **) buffer) = user_buffer;
	*size = user_buffer_size;
	user_buffer = NULL;
	user_buffer_size = 0;
	return MPI_SUCCESS;
}

int MPI_Isend(void *buf, int count, MPI_Datatype datatype, int dest,
	int tag, MPI_Comm comm, MPI_Request *request)
{
	Comm *real_comm;
	Group *real_group;
	int dest_world_rank;
	int ierr;
	
	real_comm = OOM_comm_lookup(comm);
	if (real_comm == NULL) return MPI_ERR_COMM;
	real_group = OOM_group_lookup(real_comm->group);
	if (real_group == NULL) return MPI_ERR_COMM;

	if (dest < 0 || dest > real_group->num_members) return MPI_ERR_RANK;
	dest_world_rank = real_group->members[dest];
		
	if (count < 0) return MPI_ERR_COUNT;
	
	/* check for valid datatype and that datatype is commited,
	   return MPI_ERR_TYPE if not */
	
	if (tag < 0) return MPI_ERR_TAG;

	ierr = CM_Send(buf, count, datatype, real_comm->source, dest,
		real_comm->p2p_context, tag, dest_world_rank, request);
	if (ierr != CM_SUCCESS) return MPI_ERR_INTERN;
	
	return MPI_SUCCESS;
}

int MPI_Ibsend(void *buf, int count, MPI_Datatype datatype, int dest,
	int tag, MPI_Comm comm, MPI_Request *request)
{
	return MPI_ERR_OTHER;
}

int MPI_Issend(void *buf, int count, MPI_Datatype datatype, int dest,
	int tag, MPI_Comm comm, MPI_Request *request)
{
	return MPI_Isend(buf, count, datatype, dest, tag, comm, request);
}

int MPI_Irsend(void *buf, int count, MPI_Datatype datatype, int dest,
	int tag, MPI_Comm comm, MPI_Request *request)
{
	return MPI_Isend(buf, count, datatype, dest, tag, comm, request);
}

int MPI_Irecv(void *buf, int count, MPI_Datatype datatype, int source,
	int tag, MPI_Comm comm, MPI_Request *request)
{
	Comm *real_comm;
	Group *real_group;
	int dest, world_rank, ierr;
	
	real_comm = OOM_comm_lookup(comm);
	if (real_comm == NULL) return MPI_ERR_COMM;
	real_group = OOM_group_lookup(real_comm->group);
	if (real_group == NULL) return MPI_ERR_COMM;
	dest = 0;
	world_rank = get_world_rank();
	while (dest < real_group->num_members)
	{
		if (real_group->members[dest] == world_rank) break;
		dest++;
	}
	if (dest == real_group->num_members) return MPI_ERR_COMM;
	
	if (count < 0) return MPI_ERR_COUNT;
	
	/* check for valid datatype and that datatype is commited,
	   return MPI_ERR_TYPE if not */
	
	if (tag < 0 && tag != MPI_ANY_TAG)
		return MPI_ERR_TAG;
	
	if ((source < 0 || source >= real_group->num_members)
		&& source != MPI_ANY_SOURCE) return MPI_ERR_RANK;

	ierr = CM_Recv(buf, count, datatype, source, dest, real_comm->p2p_context,
		tag, request);
	if (ierr != CM_SUCCESS) return MPI_ERR_INTERN;
	
	return MPI_SUCCESS;
}

int MPI_Wait(MPI_Request *request, MPI_Status *status)
{
	Request *real_request;
	int ierr;

	if (*request == MPI_REQUEST_NULL)
	{
		set_status_empty(status);
		return MPI_SUCCESS;
	}
	
	real_request = OOM_request_lookup(*request);
	if (real_request == NULL) return MPI_ERR_REQUEST;
	
	if (real_request->status == INACTIVE)
	{
		set_status_empty(status);		
		return MPI_SUCCESS;
	}
	
	ierr = CM_Wait2(real_request, status);
	if (ierr != CM_SUCCESS) return MPI_ERR_OTHER;
	
	if (! real_request->is_persistent)
	{
		*request = MPI_REQUEST_NULL;
	}
	
	return MPI_SUCCESS;
}

int MPI_Test(MPI_Request *request, int *flag, MPI_Status *status)
{
	Request *real_request;
	int ierr;

	if (*request == MPI_REQUEST_NULL)
	{
		set_status_empty(status);
		*flag = 1;
		return MPI_SUCCESS;
	}
	
	real_request = OOM_request_lookup(*request);
	if (real_request == NULL) return MPI_ERR_REQUEST;
	
	if (real_request->status == INACTIVE)
	{
		set_status_empty(status);
		*flag = 1;		
		return MPI_SUCCESS;
	}
	
	ierr = CM_Test2(real_request, flag, status);
	if (ierr != CM_SUCCESS) return MPI_ERR_OTHER;
	
	if (*flag && ! real_request->is_persistent)
	{
		*request = MPI_REQUEST_NULL;
	}
	
	return MPI_SUCCESS;
}

int MPI_Request_free(MPI_Request *request)
{
	if (*request == MPI_REQUEST_NULL) return MPI_ERR_REQUEST;
	
	*request = MPI_REQUEST_NULL;
	
	return MPI_SUCCESS;
}

int MPI_Waitany(int count, MPI_Request *array_of_requests, int *index,
	MPI_Status *status)
{
	apr_pool_t *mp;
	Request **real_requests;
	int i, active_requests, retval;
	
	retval = MPI_SUCCESS;
	
	apr_pool_create(&mp, NULL);
	real_requests = (Request **) apr_palloc(mp, sizeof(Request *) * count);
	
	active_requests = 0;
	for (i = 0; i < count; i++)
	{
		if (array_of_requests[i] == MPI_REQUEST_NULL)
		{
			real_requests[i] = NULL;
		}
		else
		{
			real_requests[i] = OOM_request_lookup(array_of_requests[i]);
			if (real_requests[i] == NULL)
			{
				retval = MPI_ERR_REQUEST;
				break;
			}
			else if (real_requests[i]->status == INACTIVE)
			{
				real_requests[i] = NULL;
			}
			else
			{
				active_requests++;
			}
		}
	}
	
	if (retval == MPI_SUCCESS)
	{
		if (active_requests == 0)
		{
			*index = MPI_UNDEFINED;
			set_status_empty(status);
		}
		else
		{
			CM_Waitany(count, real_requests, index, status);
		}
	}
	
	apr_pool_destroy(mp);
	
	return retval;
}

int MPI_Testany(int count, MPI_Request *array_of_requests, int *index,
	int *flag, MPI_Status *status)
{
	int i;
	*index = MPI_UNDEFINED;
	for (i = 0; i < count; i++)
	{
		MPI_Test(array_of_requests + i, flag, status);
		if (flag)
		{
			/* FIXME: should not set index if request was null or inactive */
			*index = i;
		}
	}
	return MPI_SUCCESS;
}

int MPI_Waitall(int count, MPI_Request *array_of_requests,
	MPI_Status *array_of_statuses)
{
	int i, retval, ierr;
	retval = MPI_SUCCESS;
	for (i = 0; i < count; i++)
	{
		array_of_statuses[i].MPI_ERROR = MPI_PENDING;
		ierr = MPI_Wait(array_of_requests + i, array_of_statuses + i);
		array_of_statuses[i].MPI_ERROR = ierr;
		if (ierr != MPI_SUCCESS)
		{
			retval = MPI_ERR_IN_STATUS;
		}
	}
	return retval;
}

int MPI_Testall(int count, MPI_Request *array_of_requests, int *flag,
	MPI_Status *array_of_statuses)
{
	int i, retval, ierr, iflag;
	retval = MPI_SUCCESS;
	for (i = 0; i < count; i++)
	{
		array_of_statuses[i].MPI_ERROR = MPI_PENDING;
		ierr = MPI_Test(array_of_requests + i, &iflag, array_of_statuses + i);
		*flag = *flag && iflag;
		array_of_statuses[i].MPI_ERROR = ierr;
		if (ierr != MPI_SUCCESS)
		{
			retval = MPI_ERR_IN_STATUS;
		}
	}
	return retval;
}

int MPI_Waitsome(int incount, MPI_Request *array_of_requests, int *outcount,
	int *array_of_indeces, MPI_Status *array_of_statuses)
{
	apr_pool_t *mp;
	Request **real_requests;
	int i, active_requests, retval;
	
	retval = MPI_SUCCESS;
	
	apr_pool_create(&mp, NULL);
	real_requests = (Request **) apr_palloc(mp, sizeof(Request *) * incount);
	
	active_requests = 0;
	for (i = 0; i < incount; i++)
	{
		if (array_of_requests[i] == MPI_REQUEST_NULL)
		{
			real_requests[i] = NULL;
		}
		else
		{
			real_requests[i] = OOM_request_lookup(array_of_requests[i]);
			if (real_requests[i] == NULL)
			{
				retval = MPI_ERR_REQUEST;
				break;
			}
			else if (real_requests[i]->status == INACTIVE)
			{
				real_requests[i] = NULL;
			}
			else
			{
				active_requests++;
			}
		}
	}
	
	if (retval == MPI_SUCCESS)
	{
		if (active_requests == 0)
		{
			*outcount = MPI_UNDEFINED;
		}
		else
		{
			fprintf(stderr,"Going into CM_Waitsome.\n");
			CM_Waitsome(incount, real_requests, outcount, array_of_indeces,
				array_of_statuses);
			fprintf(stderr,"Leaving CM_Waitsome.\n");
		}
	}
	
	apr_pool_destroy(mp);
	
	return retval;
}

int MPI_Testsome(int incount, MPI_Request *array_of_requests, int *outcount,
	int *array_of_indices, MPI_Status *array_of_statuses)
{
	Request *real_request;
	int i, active_requests, retval, flag, source, tag, error;
	
	retval = MPI_SUCCESS;
		
	active_requests = 0;
	*outcount = 0;
	for (i = 0; i < incount; i++)
	{
		if (array_of_requests[i] == MPI_REQUEST_NULL)
		{
			continue;
		}
		else
		{
			real_request = OOM_request_lookup(array_of_requests[i]);
			if (real_request == NULL)
			{
				retval = MPI_ERR_REQUEST;
				break;
			}
			else if (real_request->status == INACTIVE)
			{
				continue;
			}
			else
			{
				active_requests++;
				CM_Test2(real_request, &flag, array_of_statuses + *outcount);
				if (flag)
				{
					array_of_indices[*outcount] = i;
					*outcount++;
				}
			}
		}
	}
	
	if (retval == MPI_SUCCESS && active_requests == 0)
	{
		*outcount = MPI_UNDEFINED;
	}
	
	return retval;
}

int MPI_Iprobe(int source, int tag, MPI_Comm comm, int *flag,
	MPI_Status *status)
{
	Comm *real_comm;
	
	real_comm = OOM_comm_lookup(comm);
	if (real_comm == NULL) return MPI_ERR_COMM;
	
	CM_Iprobe(source, real_comm->source, real_comm->p2p_context, tag, flag,
		status);
	
	return MPI_SUCCESS;
}

int MPI_Probe(int source, int tag, MPI_Comm comm, MPI_Status *status)
{
	Comm *real_comm;
	
	real_comm = OOM_comm_lookup(comm);
	if (real_comm == NULL) return MPI_ERR_COMM;
	
	CM_Probe(source, real_comm->source, real_comm->p2p_context, tag, status);
	
	return MPI_SUCCESS;
}

int MPI_Cancel(MPI_Request *request)
{
	return MPI_SUCCESS;
}

int MPI_Test_cancelled(MPI_Status *status, int *flag)
{
	*flag = 0;
	
	return MPI_SUCCESS;
}

int MPI_Send_init(void *buf, int count, MPI_Datatype datatype, int dest,
	int tag, MPI_Comm comm, MPI_Request *request)
{
	Comm *real_comm;
	Group *real_group;
	Request *real_request;
	
	real_comm = OOM_comm_lookup(comm);
	if (real_comm == NULL) return MPI_ERR_COMM;
	
	real_group = OOM_group_lookup(real_comm->group);
	if (real_group == NULL) return MPI_ERR_COMM;
	
	*request = OOM_request_create();
	real_request = OOM_request_lookup(*request);
	if (real_request == NULL) return MPI_ERR_INTERN;
	
	real_request->req_type = SEND;
	real_request->buf = buf;
	real_request->count = count;
	real_request->datatype = datatype;
	real_request->request_env.source = real_comm->source;
	real_request->request_env.dest = dest;
	real_request->request_env.context = real_comm->p2p_context;
	real_request->request_env.tag = tag;
	real_request->dest_world_rank = real_group->members[dest];
	MPI_Type_size(datatype, &(real_request->num_bytes));
	real_request->num_bytes *= count;
	real_request->status = INACTIVE;
	real_request->is_persistent = 1;
	
	return MPI_SUCCESS;
}

int MPI_Bsend_init(void *buf, int count, MPI_Datatype datatype, int dest, int tag, MPI_Comm comm, MPI_Request *request)
{
	return MPI_SUCCESS;
}

int MPI_Ssend_init(void *buf, int count, MPI_Datatype datatype, int dest, int tag, MPI_Comm comm, MPI_Request *request)
{
	return MPI_Send_init(buf, count, datatype, dest, tag, comm, request);
}

int MPI_Rsend_init(void *buf, int count, MPI_Datatype datatype, int dest, int tag, MPI_Comm comm, MPI_Request *request)
{
	return MPI_Send_init(buf, count, datatype, dest, tag, comm, request);
}

int MPI_Recv_init(void *buf, int count, MPI_Datatype datatype, int source,
	int tag, MPI_Comm comm, MPI_Request *request)
{
	Comm *real_comm;
	Request *real_request;
	
	real_comm = OOM_comm_lookup(comm);
	if (real_comm == NULL) return MPI_ERR_COMM;
	
	*request = OOM_request_create();
	real_request = OOM_request_lookup(*request);
	if (real_request == NULL) return MPI_ERR_INTERN;
	
	real_request->req_type = RECV;
	real_request->buf = buf;
	real_request->count = count;
	real_request->datatype = datatype;
	real_request->request_env.source = source;
	real_request->request_env.dest = real_comm->source;
	real_request->request_env.context = real_comm->p2p_context;
	real_request->request_env.tag = tag;
	real_request->status = INACTIVE;
	real_request->is_persistent = 1;
	
	return MPI_SUCCESS;
}

int MPI_Start(MPI_Request *request)
{
	Request *real_request;
	
	real_request = OOM_request_lookup(*request);
	if (real_request == NULL || ! real_request->is_persistent)
	{
		return MPI_ERR_REQUEST;
	}
	
	CM_Start(real_request);
	
	return MPI_SUCCESS;
}

int MPI_Startall(int num_requests, MPI_Request *array_of_requests)
{
	int i;
	
	for (i = 0; i < num_requests; i++)
	{
		MPI_Start(array_of_requests[i]);
	}
	
	return MPI_SUCCESS;
}

/****************************************************
 * \author James Edmondson
 * \brief Sends from a provided send buffer and receives into a provided
 *        receive buffer. Both are executed asynchronously with our
 *        communication manager.
 *****************************************************/

int MPI_Sendrecv(void *sendbuf, int sendcount, MPI_Datatype sendtype,
	int dest, int sendtag, void *recvbuf, int recvcount,
	MPI_Datatype recvtype, int source, int recvtag, MPI_Comm comm,
	MPI_Status *status)
{
	Comm *real_comm = NULL;
	Group *real_group = NULL;
	Request * real_request = NULL;
	MPI_Request requests[2];
	int my_group_rank;
	int tsource;
	int ierr;
	int tag;
	real_comm = OOM_comm_lookup(comm);
	if (real_comm == NULL) return MPI_ERR_COMM;
	real_group = OOM_group_lookup(real_comm->group);
	if (real_group == NULL) return MPI_ERR_COMM;
	
	my_group_rank = real_comm->source;

	/* perform send and recv asynchronously */
	
	CM_Send(sendbuf,sendcount,sendtype,my_group_rank,dest,
			real_comm->p2p_context,sendtag,
			real_group->members[dest],&requests[0]);
	CM_Recv(recvbuf,recvcount,recvtype,source,my_group_rank,
			real_comm->p2p_context,recvtag,&requests[1]);

	/* but check the send completion first */
	
	real_request = OOM_request_lookup(requests[0]);
	CM_Wait(real_request,&tsource,&tag,&ierr);

	/* return status of the wait on the receive  */
	
	return MPI_Wait(&requests[1], status);
}

/*****************************************************
 * \author James Edmondson
 * \brief Sends and receives into a buffer in the same operation.
 *        we currently handle this by allocating a temporary buffer
 *        and memcpying back into the provided buffer. HOWEVER, this
 *        should not work with custom user types. We should be using
 *        MPI_PACK.
 *****************************************************/

int MPI_Sendrecv_replace(void *buf, int count, MPI_Datatype datatype,
	int dest, int sendtag, int source, int recvtag, MPI_Comm comm,
	MPI_Status *status)
{
	Comm *real_comm = NULL;
	Group *real_group = NULL;
	Request * real_request = NULL;
	MPI_Request requests[2];
	int my_group_rank;
	int tsource;
	int ierr;
	int tag;
	int datasize;
	apr_pool_t *mp;
	real_comm = OOM_comm_lookup(comm);
	if (real_comm == NULL) return MPI_ERR_COMM;
	real_group = OOM_group_lookup(real_comm->group);
	if (real_group == NULL) return MPI_ERR_COMM;

	ierr = MPI_Type_extent(datatype,&datasize);
	if( ierr != MPI_SUCCESS ) return MPI_ERR_TYPE;
	
	apr_pool_create(&mp, NULL);
	void * recvbuf = (void *) apr_palloc(mp, datasize * count);
	
	my_group_rank = real_comm->source;

	/* perform send first */
	
	CM_Send(buf,count,datatype,my_group_rank,dest,
			real_comm->p2p_context,sendtag,
			real_group->members[dest],&requests[0]);

	CM_Recv(recvbuf,count,datatype,source,my_group_rank,
			real_comm->p2p_context,recvtag,&requests[1]);
	/* wait on the send to complete */
	
	real_request = OOM_request_lookup(requests[0]);
	CM_Wait(real_request,&tsource,&tag,&ierr);
	
	/* return status of the wait on the receive  */
	
	ierr = MPI_Wait(&requests[1], status);

	memcpy(recvbuf,buf,datasize * count);

	return ierr;
}
